# websocket_server.py

import base64
import json
from fastapi import WebSocket, WebSocketDisconnect, APIRouter
from thirdpart.xf_iat.iflytek_ws_client import IflytekClient
import asyncio
from settings import settings


iflytek_ws_stt_router = APIRouter()

@iflytek_ws_stt_router.websocket("/ws/speech", name ="语音识别WebSocket接口")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    buffer = bytearray()
    loop = asyncio.get_event_loop()

    def on_result(result: str):
        print(result)
        asyncio.run_coroutine_threadsafe(websocket.send_text(result), loop)

 
    client = IflytekClient(
    appid=settings.IFLYTEK_APPID,
    apikey=settings.IFLYTEK_APIKEY,
    apisecret=settings.IFLYTEK_APISECRET,
    on_result_callback=on_result,
    )
  

    client.connect()

    try:
        while True:
            audio_chunk = await websocket.receive_bytes()
            buffer.extend(audio_chunk)
            while len(buffer) >= 8000:
                client.send_audio(buffer[:8000])
                buffer = buffer[8000:]
    except WebSocketDisconnect:
        print("WebSocket disconnected")
        client.close()

# --------- 新增 /ws/speech_v2 接口，支持心跳 + base64 音频帧 --------- #
@iflytek_ws_stt_router.websocket("/ws/speech_v2", name="语音识别WebSocket接口_v2")
async def websocket_endpoint_v2(websocket: WebSocket):
    """
    新版本 WebSocket 接口，客户端发送 JSON：
      {
        "type": "heartbeat"
      }
    或
      {
        "type": "audio",
        "frame_type": 0|1|2,          # 0x00 起始帧, 0x01 中间帧, 0x02 结束帧
        "data": "<base64_encoded_audio_chunk>"
      }
    心跳包收到后立即回 "pong"，音频数据收到后才初始化 IflytekClient 并转发解码发送。
    """
    await websocket.accept()
    print("[v2] WebSocket 连接已接受")
    
    buffer = bytearray()
    client = None
    loop = asyncio.get_event_loop()

    import re

    def on_result_v2(recognized_text: str):
        print(f"[v2][Iflytek] 识别结果回调: {recognized_text}")

        # 去掉前后空白
        cleaned_text = recognized_text.strip()

        # 如果内容为空，或只包含标点符号，跳过发送
        if not cleaned_text or re.fullmatch(r"[^\w\u4e00-\u9fff]+", cleaned_text):
            print("[v2][Iflytek] 空或纯标点，跳过发送")
            return

        # 正常发送
        asyncio.run_coroutine_threadsafe(
            websocket.send_text(json.dumps({
                "type": "result",
                "text": recognized_text
            })),
            loop
        )


    # # 定义 Iflytek 识别结果回调
    # def on_result_v2(recognized_text: str):
    #     print(f"[v2][Iflytek] 识别结果回调: {recognized_text}")

    #     if recognized_text.strip():  # 非空时才发送
    #         asyncio.run_coroutine_threadsafe(
    #             websocket.send_text(json.dumps({
    #                 "type": "result",
    #                 "text": recognized_text
    #             })), 
    #             loop
    #         )
    #     else:
    #         print("[v2][Iflytek] 空识别结果，跳过发送")

    # # 定义 Iflytek 识别结果回调
    # def on_result_v2(recognized_text: str):
    #     print(f"[v2][Iflytek] 识别结果回调: {recognized_text}")
    #     # 异步将识别结果发给客户端
    #     asyncio.run_coroutine_threadsafe(websocket.send_text(json.dumps({
    #         "type": "result",
    #         "text": recognized_text
    #     })), loop)

    try:
        while True:
            # 收到客户端 JSON 文本
            raw_msg = await websocket.receive_text()
            #print(f"[v2] 收到客户端消息: {raw_msg}")
            
            try:
                msg_obj = json.loads(raw_msg)
            except json.JSONDecodeError as e:
                print(f"[v2] JSON 解析失败: {e}，跳过该消息")
                continue

            msg_type = msg_obj.get("type")
            
            # ---- 处理心跳 ----
            if msg_type == "heartbeat":
                print("[v2] 收到心跳包，回复 pong")
                await websocket.send_text(json.dumps({"type": "pong"}))
                continue

            # ---- 处理音频数据 ----
            elif msg_type == "audio":
                frame_type = msg_obj.get("frame_type")
                b64_data = msg_obj.get("data", "")
                print(f"[v2] 收到音频帧，frame_type={frame_type}, base64 长度={len(b64_data)}")

                # base64 解码
                try:
                    audio_bytes = base64.b64decode(b64_data)
                except Exception as e:
                    print(f"[v2] base64 解码失败: {e}, 跳过此帧")
                    continue

                # 如果是起始帧（0），则初始化 IflytekClient 并 connect
                if frame_type == 0:
                    print("[v2] 起始帧到达，初始化 IflytekClient 并建立连接")
                    # 如果之前已有 client 实例，先关闭
                    # if client is not None:
                    #     print("[v2] 之前的 client 未关闭，先执行 close()")
                    #     client.close()
                    #     client = None
                    buffer.clear()

                    client = IflytekClient(
                        appid=settings.IFLYTEK_APPID,
                        apikey=settings.IFLYTEK_APIKEY,
                        apisecret=settings.IFLYTEK_APISECRET,
                        on_result_callback=on_result_v2,
                    )
                    client.connect()
                    # print("[v2] IflytekClient.connect() 已调用")

                # 将解码后音频追加到 buffer
                buffer.extend(audio_bytes)
                # print(f"[v2] 累计 buffer 长度={len(buffer)} 字节")

                # 按 8000 字节一块发送给 Iflytek
                while len(buffer) >= 1280 and client is not None:
                    chunk = bytes(buffer[:1280])
                    # print("[v2] 发送 1280 字节到 IflytekClient.send_audio()")
                    client.send_audio(chunk)
                    buffer = buffer[1280:]
                    # print(f"[v2] 剩余 buffer 长度={len(buffer)} 字节")

                # 如果是结束帧（2），将剩余 buffer 全部发送，然后关闭 client
                if frame_type == 2 and client is not None:
                    if len(buffer) > 0:
                        print(f"[v2] 结束帧到达，发送剩余 {len(buffer)} 字节到 IflytekClient.send_audio()")
                        client.send_audio(bytes(buffer))
                        buffer.clear()
                    print("[v2] 结束帧已处理，关闭 IflytekClient")
                    client.close()
                    client = None
                    print("[v2] IflytekClient 已关闭")
                    # 继续等待下一轮连接或直接退出循环也行，根据业务决定
                    # 这里选择继续循环，以便客户端可再次发送新的起始帧开启下一次识别
                    continue

            else:
                print(f"[v2] 未知消息类型: {msg_type}, 跳过处理")
                continue

    except WebSocketDisconnect:
        print("[v2] 客户端 WebSocket 已断开")
        if client is not None:
            print("[v2] 断开时关闭 IflytekClient")
            client.close()